import os
import shutil
from threading import Lock
import requests
import threading
import multiprocessing

class CallbackDict:
    def put(a1=False, a2=0):pass
    def get(a1, a2=False, a3=0):pass

class TDownload:
    def __init__(self, interface):
        self.__requests = requests
        self.__url = None
        self.__file = None
        self.__lock = Lock()
        self.th_running = {}
        self.th_running_num = 0
        self.interface = interface

    def __download_operation(self, start, end, another_origin):
        headers = {
            'Range': f'bytes={start}-{end}',
        }
        try:
            r = self.__requests.get(self.__url, stream=True, headers=headers)  # stream=True设置流式下载
        except Exception as ept:
            print("download failed", self.__url, '"', ept)
            if another_origin:
                r = self.__requests.get(self.__url, stream=True, headers=headers)  # stream=True设置流式下载
            else:
                exit()
        # 分多次下载数据
        pos = start
        for i in r.iter_content(chunk_size=1024):  # 设置每次获取的大小
            if i:  # 判断是否为空数据
                self.__lock.acquire()
                self.__file.seek(pos)
                self.__file.write(i)
                self.__lock.release()

                pos += 1024

    def download_main(self, url, path, thread_num, size, callback, another_origin):
        self.__url = url
        if os.path.isfile(path):
            return
        try:
            os.makedirs(os.path.dirname(path))
        except Exception:
            pass
        self.__file = open(path, "wb")
        th = []
        # point_list = []
        # remain_size = size - (size % thread_num)
        # size -= remain_size
        # start = 0
        # for i in range(thread_num - 1):
        #     end = start + (size // thread_num)
        #     point_list.append((start, end))
        #     start = end + 1
        # point_list.append((start, start + remain_size))
        point_list = self.cut(size, thread_num)
        for p in point_list:
            th.append(threading.Thread(target=self.__download_operation, args=(p[0], p[1], another_origin)))
        for i in th:
            i.setDaemon(True)
            i.start()
        for i in th:
            i.join()
        self.__file.close()
    
    def cut(self, filesize, chuck):
        step = filesize // chuck
        arr = list(range(0, filesize, step))
        result = []
        for i in range(len(arr)-1):
            s_pos, e_pos = arr[i], arr[i+1]-1
            result.append([s_pos, e_pos])
        result[-1][-1] = filesize-1
        return result

    def download_file(self, url, path, thread_num, size, callback = CallbackDict, another_origin = False):
        t1 = threading.Thread(target=self.download_main,
                                   args=(url, path, thread_num, size, callback, another_origin))
        self.th_running_num += 1
        self.th_running.update({self.th_running_num: t1})
        self.th_running[self.th_running_num].run()

    def get_info(self, url):
        ret = self.__requests.get(url)
        ret = ret.text
        return ret

class PDownload:
    def __init__(self, interface):
        self.use_process = 0
        self.interface = interface
    
    def download_main(self, url_list, path_list, thread_num, size_list, callback):
        for i in range(len(url_list)):
            url, path, size = url_list[i], path_list[i], size_list[i]
            self.interface("Download.TDownload").download_file(url, path, thread_num, size, callback)
            callback.put("[download] " + url, True)

    def download_file(self, url_list, path_list, thread_num, process_num, size_list, callbak = CallbackDict):
        use_list = []
        use_path_list = []
        use_size_list = []
        for i in range(process_num):
            use_list.append([])
            use_path_list.append([])
            use_size_list.append([])
        for i in range(len(url_list)):
            use_list[i % process_num].append(url_list[i])
            use_path_list[i % process_num].append(path_list[i])
            use_size_list[i % process_num].append(size_list[i])
        for i in range(process_num):
            p1 = multiprocessing.Process(target=self.download_main, args=(use_list[i], use_path_list[i], thread_num, use_size_list[i], callbak))
            p1.start()

class Download:
    def __init__(self, interface) -> None:pass

    def download_file(self, url, path):
        with open(path, "wb") as f, requests.get(url, stream=True) as res:
            shutil.copyfileobj(res.raw, f)